/*
 * EquilibriumSolver.cpp
 *
 *  Created on: 26 Jul 2011
 *      Author: Allan
 */

#include "EquilibriumSolver.h"

// C++ includes
#include <iomanip>
using namespace std;

// Boost includes
#include <boost/foreach.hpp>
using namespace boost;

// GeoReact includes
#include <Core/EquilibriumConditions.h>
#include <Core/EquilibriumProblem.h>

EquilibriumSolver::EquilibriumSolver() :
phase_threshold(1.0E-14), tolerance(1.0E-6), maxiterations(200)
{}

EquilibriumSolver::EquilibriumSolver(const Multiphase& multiphase, const ReactionSystem& reactions) :
phase_threshold(1.0E-14), tolerance(1.0E-6), maxiterations(200)
{
	Initialiaze(multiphase, reactions);
}

void EquilibriumSolver::Initialiaze(const Multiphase& multiphase, const ReactionSystem& equilibrium_reactions)
{
	// Set data members
	this->multiphase = multiphase;
	this->equilibrium_reactions = equilibrium_reactions;
	
	// Initialiaze the stoichiometric matrix of the system of equilibrium-controlled reactions
	uei = equilibrium_reactions.AssembleStoichiometricMatrix(multiphase.GetSpecies());
	
	// Initialiaze the Newton solver
	newton_solver.Initialiaze(multiphase.GetNumSpecies());
	
	InitialiazeMappingReactionToPhase();
}

void EquilibriumSolver::SetConvergenceOptions(double tolerance, uint maxiterations)
{
	this->tolerance = tolerance;
	this->maxiterations = maxiterations;
}

void AppplyUnderRelaxation(const VectorXd& nOld, VectorXd& n)
{
	for(unsigned i = 0; i < unsigned(n.rows()); ++i)
	{
		const double delta  = n[i] - nOld[i];
		const double factor = max(1.0, -2.0 * delta/nOld[i]);
		
		n[i] = nOld[i] + delta/factor;
	}
}

ConvergenceDiagnostics EquilibriumSolver::Solve(const EquilibriumConditions& econditions, double T, double P, VectorXd& n)
{
	// Auxiliary vector variables 
	VectorXd nOld, a, ln_a, ln_Ke, hj;
	
	// Auxiliary matrix variables
	MatrixXd Aii;
	
	// Calculate the natural logarithm of the equilibrium constants of the equilibrium-controlled reactions
	ln_Ke = equilibrium_reactions.EquilibriumConstants(T, P).array().log();
	
	// The equilibrium condition matrix
	const MatrixXd& alpha = econditions.GetConditionMatrix();
	
	// The current state of the convergence
	ConvergenceDiagnostics diagnostics = {0, false};
	
	while(!diagnostics.converged && diagnostics.num_iterations++ < maxiterations)
	{
		// Calculate the activities of all species
		a = multiphase.Activities(T, P, n);
		
		// Calculate the natural logarithm of the activity derivative of all species
		Aii = a.asDiagonal().inverse() * multiphase.ActivityDerivatives(T, P, n, a);
		
		// Calculate the natural logarithm of the activities
		ln_a = a.array().log();
		
		// Calculate the (hj) vector with the equilibrium condition functions
		hj = econditions.EvaluateConditionFunctions(T, P, n, a);
		
		// Create an equilibrium problem instance to calculate the function F and jacobian J
		EquilibriumProblem problem(Aii, uei, alpha, hj, ln_a, ln_Ke);
		
		// Update the previous composition state (nOld)
		nOld = n;
		
		// Perform one iteration with the Newton's method and retrieve the convergence state
		newton_solver.Iterate(problem, n);
		
		// Ensure non-negative values of compositions
		AppplyUnderRelaxation(nOld, n);
		
		diagnostics.converged = CheckConvergence(n);
	}
	
	return diagnostics;
}

void EquilibriumSolver::InitialiazeMappingReactionToPhase()
{
	reaction_to_phase.clear();
	
	BOOST_FOREACH(const Reaction& reaction, equilibrium_reactions.GetReactions())
	{
		set<Index> phase_indexes;
		
		BOOST_FOREACH(const string& species, reaction.GetSpecies())
			phase_indexes.insert(multiphase.PhaseIndex(species));
		
		reaction_to_phase.push_back(phase_indexes);
	}
}

bool EquilibriumSolver::CheckConvergence(const VectorXd& n) const
{
	const VectorXd& F = newton_solver.GetResidualFunction();
	
	const unsigned N  = multiphase.GetNumSpecies();
	const unsigned Ne = equilibrium_reactions.GetNumReactions();
	const unsigned Nj = N - Ne;
	
	VectorXd n_total_phases = multiphase.TotalMolesPhases(n);
	
	const double n_total = n_total_phases.sum();
	
	set<Index> nonexistent_phases;
	
	for(unsigned iphase = 0; iphase < multiphase.GetNumPhases(); ++iphase)
		if(n_total_phases[iphase] < phase_threshold * n_total)
			nonexistent_phases.insert(iphase);
	
	double residual_norm = F.segment(0, Nj).array().abs().sum();
	
	for(unsigned ireaction = 0; ireaction < equilibrium_reactions.GetNumReactions(); ++ireaction)
	{
		bool discard_reaction = false;
		
		BOOST_FOREACH(Index iphase, reaction_to_phase[ireaction])
			if(nonexistent_phases.count(iphase))
				{ discard_reaction = true; break; }
		
		if(!discard_reaction)
			residual_norm += abs(F[Nj + ireaction]);
	}
	
	return residual_norm < tolerance;
}
